package com.moss.cloud.gateway.filter;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.moss.cloud.common.core.constant.Constant;
import com.moss.cloud.common.core.model.Result;
import com.moss.cloud.common.core.utils.Sm4Util;
import com.moss.cloud.common.core.utils.WebFluxUtil;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.reactivestreams.Publisher;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.*;

/**
 * 自定义密文加解密
 *
 * @author 瑾年
 * @date 2024年1月25日
 */
@Slf4j
@Component
public class CustomCryptoFilter implements GlobalFilter, Ordered {
    /**
     * 统一返回字段
     */
    private static final String DATA = "data";
    private static final String CODE = "code";
    private static final String JSON_TYPE = "json";
    /**
     * 定义拦截加密请求方式
     */
    private static final Set<HttpMethod> NEED_FILTER_METHOD_SET = new HashSet<>(Arrays.asList(HttpMethod.GET, HttpMethod.POST, HttpMethod.PUT));

    /**
     * 定义拦截加密请求ContentType
     * 知道为啥只拦截APPLICATION_JSON的请求吗？
     * 因为文件下载，上传之类的请求没办法加密，加密之后的文件流也很难解析
     */
    private static final Set<MediaType> NEED_FILTER_MEDIA_TYPE_SET = new HashSet<>(Collections.singletonList(MediaType.APPLICATION_JSON));


    @Value("${gateway.crypto.enabled}")
    private boolean enabled;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        // 获取请求体
        ServerHttpRequest request = exchange.getRequest();
        // 获取响应体
        ServerHttpResponse response = exchange.getResponse();
        // 请求头
        HttpHeaders headers = request.getHeaders();
        // 请求方法
        HttpMethod method = request.getMethod();
        if (Boolean.FALSE.equals(enabled)) {
            log.info("请求参数加解密未开启，跳过：{}", false);
            return chain.filter(exchange);
        }
        // 满足条件，进行过滤
        if (isNeedFilterMethod(method) && isNeedFilterContentType(headers.getContentType())) {
            return DataBufferUtils.join(request.getBody())
                    .flatMap(dataBuffer -> {
                        try {
                            // 获取请求参数
                            String originalRequestBody = getOriginalRequestBody(dataBuffer);

                            // 解密请求参数
                            String decryptRequestBody = decryptRequest(originalRequestBody);

                            // 装饰新的请求体
                            ServerHttpRequestDecorator requestDecorator = serverHttpRequestDecorator(request, decryptRequestBody);

                            // 装饰新的响应体
                            ServerHttpResponseDecorator responseDecorator = serverHttpResponseDecorator(response);

                            // 使用新的请求和响应转发
                            ServerWebExchange serverWebExchange = exchange.mutate().request(requestDecorator).response(responseDecorator).build();

                            // 放行拦截
                            return chain.filter(serverWebExchange);
                        } catch (Exception e) {
                            log.error("密文过滤器加解密错误", e);
                            return WebFluxUtil.cryptoError(exchange);
                        } finally {
                            DataBufferUtils.release(dataBuffer);
                        }
                    });
        }
        return chain.filter(exchange);
    }

    private String decryptRequest(String originalRequestBody) {
        log.info("请求参数解密，原文：{}", originalRequestBody);
        JSONObject json = JSONObject.parseObject(originalRequestBody);
        String enData = json.getString("enData");
        String decrypted = Sm4Util.decrypt("1234567891011122", enData);
        log.info("请求参数解密，明文：{}", decrypted);
        return decrypted;
    }

    private String encryptResponse(String originalResponseBody) {
        JSONObject result = JSONObject.parseObject(originalResponseBody);
        // 只对data字段进行加密处理
        Object data = result.get(DATA);
        if (Objects.nonNull(data) && result.getString(CODE).equals(Constant.SYSTEM_SUCCESS)) {
            // 秘钥需要动态生成，动态解析，这里是为了测试
            String newData = Sm4Util.encrypt("1234567891011122", JSON.toJSONString(data));
            Result<?> newRes = Result.success(newData);
            log.info("响应结果加密，原文：{}", originalResponseBody);
            String response = JSONObject.toJSONString(newRes);
            log.info("响应结果加密，密文：{}", response);
            return response;
        }
        return originalResponseBody;

    }

    /**
     * 获取原始的请求参数
     *
     * @param dataBuffer 数据缓冲
     * @return 原始的请求参数
     */
    private String getOriginalRequestBody(DataBuffer dataBuffer) {
        byte[] bytes = new byte[dataBuffer.readableByteCount()];
        dataBuffer.read(bytes);
        return new String(bytes, StandardCharsets.UTF_8);
    }


    private boolean isNeedFilterMethod(HttpMethod method) {
        return NEED_FILTER_METHOD_SET.contains(method);
    }

    private boolean isNeedFilterContentType(MediaType mediaType) {
        return NEED_FILTER_MEDIA_TYPE_SET.contains(mediaType) || JSON_TYPE.equals(mediaType.getSubtype());
    }

    private ServerHttpRequestDecorator serverHttpRequestDecorator(ServerHttpRequest originalRequest, String decryptRequestBody) {
        return new ServerHttpRequestDecorator(originalRequest) {
            @NotNull
            @Override
            public HttpHeaders getHeaders() {
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll(super.getHeaders());
                httpHeaders.remove(HttpHeaders.CONTENT_LENGTH);
                httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                return httpHeaders;
            }

            @NotNull
            @Override
            public Flux<DataBuffer> getBody() {
                byte[] bytes = decryptRequestBody.getBytes(StandardCharsets.UTF_8);
                return Flux.just(new DefaultDataBufferFactory().wrap(bytes));
            }
        };
    }


    private ServerHttpResponseDecorator serverHttpResponseDecorator(ServerHttpResponse originalResponse) {
        DataBufferFactory dataBufferFactory = originalResponse.bufferFactory();
        return new ServerHttpResponseDecorator(originalResponse) {
            @NotNull
            @Override
            public Mono<Void> writeWith(@NotNull Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {
                    Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                    return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                        DataBuffer join = dataBufferFactory.join(dataBuffers);
                        byte[] byteArray = new byte[join.readableByteCount()];
                        join.read(byteArray);
                        DataBufferUtils.release(join);
                        String originalResponseBody = new String(byteArray, StandardCharsets.UTF_8);
                        // 加密
                        byte[] encryptedByteArray = encryptResponse(originalResponseBody).getBytes(StandardCharsets.UTF_8);
                        originalResponse.getHeaders().setContentLength(encryptedByteArray.length);
                        originalResponse.getHeaders().setContentType(MediaType.APPLICATION_JSON);
                        return dataBufferFactory.wrap(encryptedByteArray);
                    }));
                }
                return super.writeWith(body);
            }

            @NotNull
            @Override
            public Mono<Void> writeAndFlushWith(@NotNull Publisher<? extends Publisher<? extends DataBuffer>> body) {
                return writeWith(Flux.from(body).flatMapSequential(p -> p));
            }

            @NotNull
            @Override
            public HttpHeaders getHeaders() {
                HttpHeaders headers = new HttpHeaders();
                headers.putAll(originalResponse.getHeaders());
                return headers;
            }
        };
    }

    @Override
    public int getOrder() {
        return -1;
    }
}


